#include "MessageLoop.h"

#include <atomic>
#include <condition_variable>
#include <mutex>
#include <queue>
#include <thread>
#include <unordered_set>
#include <vector>

namespace zeroplus::mq
{

    constexpr size_t kMaxTaskSize = 10000U;
    constexpr size_t kMaxTimerSize = 100U;
    constexpr uint32_t kMaxElaspedTime = 0U;

    template <typename TypeT, typename TimeT, typename ReturnT>
    static ReturnT get_current_time() noexcept
    {
        const auto &now = TypeT::now();
        const auto &duration = std::chrono::duration_cast<TimeT>(now.time_since_epoch());
        return static_cast<ReturnT>(duration.count());
    }

    // MessageLoopImpl
    struct MessageLoopImpl
    {
        using TaskTuple = std::tuple<uint32_t, uint32_t, uint32_t, MessageLoop::Callback>;
        struct PriorityCompare
        {
            bool operator()(const TaskTuple &lhs, const TaskTuple &rhs) const
            {
                return (static_cast<uint64_t>(std::get<0>(lhs)) << 32 | std::get<1>(lhs)) >
                       (static_cast<uint64_t>(std::get<0>(rhs)) << 32 | std::get<1>(rhs));
            }
        };
        using TaskQueue = std::priority_queue<TaskTuple, std::vector<TaskTuple>, PriorityCompare>;
        std::string name;
        std::atomic_bool is_running{false};
        std::atomic_bool quit_flag{false};
        std::atomic_bool is_busy{false};
        uint32_t task_seq{0};
        std::atomic<std::thread::id> thread_id;
        MessageLoop::Callback begin_callback;
        MessageLoop::Callback end_callback;
        MessageLoop::Callback idle_callback;
        std::thread thread;
        TaskQueue task_queue;
        std::unordered_set<Timer *> timer_set;
        std::mutex mtx;
        std::condition_variable cv;
    };

    // MessageLoop
    MessageLoop::MessageLoop() : impl_(std::make_unique<MessageLoopImpl>())
    {
        impl_->name = std::to_string(reinterpret_cast<std::uintptr_t>(this));
    }

    MessageLoop::~MessageLoop()
    {
        if (impl_->is_running)
        {
            LOG_WARN("MessageLoop has still running(%s).", impl_->name.c_str());
            quit();
            wait_for_quit(1000, false);
        }
        if (impl_->thread.joinable())
        {
            impl_->thread.join();
        }
        std::unique_lock lock(impl_->mtx);
        for (auto *timer : impl_->timer_set)
        {
            timer->clear();
            if (timer->is_once_type())
            {
                lock.unlock();
                delete timer;
                lock.lock();
            }
        }
    }

    void MessageLoop::set_name(const std::string &name) { impl_->name = name; }

    const std::string &MessageLoop::get_name() { return impl_->name; }

    void MessageLoop::register_begin_handler(const Callback &callback)
    {
        if (impl_->is_running)
        {
            LOG_ERROR("MessageLoop is running and cannot be registered(%s).", impl_->name.c_str());
            return;
        }
        impl_->begin_callback = callback;
    }

    void MessageLoop::register_end_handler(const Callback &callback)
    {
        if (impl_->is_running)
        {
            LOG_ERROR("MessageLoop is running and cannot be registered(%s).", impl_->name.c_str());
            return;
        }
        impl_->end_callback = callback;
    }

    void MessageLoop::register_idle_handler(const Callback &callback)
    {
        if (impl_->is_running)
        {
            LOG_ERROR("MessageLoop is running and cannot be registered(%s).", impl_->name.c_str());
            return;
        }
        impl_->idle_callback = callback;
    }

    bool MessageLoop::run()
    {
        if (impl_->is_running || impl_->thread.joinable())
        {
            LOG_WARN("MessageLoop has already run(%s).", impl_->name.c_str());
            return false;
        }
        impl_->quit_flag = false;
        impl_->is_running = true;
        do_consume();
        return true;
    }

    bool MessageLoop::async_run()
    {
        if (impl_->is_running || impl_->thread.joinable())
        {
            LOG_WARN("MessageLoop has already run(%s).", impl_->name.c_str());
            return false;
        }
        impl_->quit_flag = false;
        impl_->is_running = true;
        impl_->thread = std::thread(&MessageLoop::do_consume, this);
        return true;
    }

    bool MessageLoop::quit()
    {
        if (!impl_->is_running || impl_->quit_flag)
        {
            return false;
        }
        {
            std::lock_guard lock(impl_->mtx);
            impl_->quit_flag = true;
        }
        impl_->cv.notify_one();
        return true;
    }

    bool MessageLoop::wait_for_idle(int ms, bool check)
    {
        if (check && is_in_same_thread())
        {
            LOG_ERROR("MessageLoop wait_for_idle in work thread(%s).", impl_->name.c_str());
            return false;
        }
        std::unique_lock lock(impl_->mtx);
        const auto &predicate_func = [this]() -> bool
        { return !impl_->is_busy.load() && impl_->task_queue.empty(); };
        if (ms == Timer::kInfinite)
        {
            impl_->cv.wait(lock, predicate_func);
            return true;
        }
        return impl_->cv.wait_for(lock, std::chrono::milliseconds(ms), predicate_func);
    }

    bool MessageLoop::wait_for_quit(int ms, bool check)
    {
        if (check && is_in_same_thread())
        {
            LOG_ERROR("MessageLoop wait_for_quit in work thread(%s).", impl_->name.c_str());
            return false;
        }
        std::unique_lock lock(impl_->mtx);
        const auto &predicate_func = [this]() -> bool
        { return !impl_->is_running.load(); };
        if (ms == Timer::kInfinite)
        {
            impl_->cv.wait(lock, predicate_func);
            return true;
        }
        return impl_->cv.wait_for(lock, std::chrono::milliseconds(ms), predicate_func);
    }

    bool MessageLoop::post_task(const Callback &callback, uint16_t priority)
    {
        if (impl_->quit_flag)
        {
            return false;
        }
        bool is_full = false;
        int retry_cnt = 0;
        do
        {
            {
                std::lock_guard lock(impl_->mtx);
                is_full = impl_->task_queue.size() >= get_max_task_count();
                if MQ_LIKELY (!is_full)
                {
                    emplace_task(callback, priority);
                }
            }
            if (is_full)
            {
                if (++retry_cnt > 10)
                {
                    std::lock_guard lock(impl_->mtx);
                    emplace_task(callback, priority);
                    LOG_WARN("Task is full(%s).", impl_->name.c_str());
                    break;
                }
                std::this_thread::sleep_for(std::chrono::milliseconds(5));
            }
        } while (is_full);
        wakeup();
        return !is_full;
    }

    bool MessageLoop::wakeup()
    {
        if (!impl_->is_running)
        {
            return false;
        }
        impl_->is_busy = true;
        impl_->cv.notify_one();
        return true;
    }

    bool MessageLoop::is_running() const { return impl_->is_running; }

    bool MessageLoop::is_ready_to_quit() const { return impl_->quit_flag; }

    bool MessageLoop::is_busy() const { return impl_->is_busy; }

    size_t MessageLoop::get_task_count() const
    {
        std::lock_guard lock(impl_->mtx);
        return impl_->task_queue.size();
    }

    size_t MessageLoop::get_max_task_count() const { return kMaxTaskSize; }

    size_t MessageLoop::get_max_timer_count() const { return kMaxTimerSize; }

    uint32_t MessageLoop::get_max_elapsed_time() const { return kMaxElaspedTime; }

    bool MessageLoop::is_in_same_thread() const { return impl_->thread_id == std::this_thread::get_id(); }

    void MessageLoop::on_begin()
    {
        if (impl_->begin_callback)
        {
            impl_->begin_callback();
        }
    }

    void MessageLoop::on_end()
    {
        if (impl_->end_callback)
        {
            impl_->end_callback();
        }
    }

    void MessageLoop::on_idle()
    {
        if (impl_->idle_callback)
        {
            impl_->idle_callback();
        }
    }

    void MessageLoop::on_task_changed(const Callback &callback, uint32_t start_time)
    {
        if (start_time > 0)
        {
            uint32_t elapsed_time =
                get_current_time<std::chrono::steady_clock, std::chrono::milliseconds, uint32_t>() - start_time;
            if (elapsed_time > get_max_elapsed_time())
            {
                on_task_timeout(callback, elapsed_time);
            }
            else
            {
                callback();
            }
        }
        else
        {
            callback();
        }
    }

    void MessageLoop::on_task_timeout(const Callback &callback, uint32_t elapsed_time)
    {
        (void)callback;
        LOG_WARN("Task is timeout with %ums(%s).", elapsed_time, impl_->name.c_str());
    }

    uint64_t MessageLoop::get_current_nano_time()
    {
        return get_current_time<std::chrono::steady_clock, std::chrono::nanoseconds, uint64_t>();
    }

    bool MessageLoop::add_timer(Timer *timer)
    {
        std::lock_guard lock(impl_->mtx);
        if (impl_->timer_set.size() >= get_max_timer_count())
        {
            LOG_WARN("Timer is full(%s).", impl_->name.c_str());
            return false;
        }
        return impl_->timer_set.emplace(timer).second;
    }

    bool MessageLoop::remove_timer(Timer *timer)
    {
        std::lock_guard lock(impl_->mtx);
        return impl_->timer_set.erase(timer) != 0;
    }

    void MessageLoop::emplace_task(const Callback &callback, uint16_t priority)
    {
        uint32_t start_time = 0;
        if (get_max_elapsed_time() > 0)
        {
            start_time = get_current_time<std::chrono::steady_clock, std::chrono::milliseconds, uint32_t>();
        }
        impl_->task_queue.emplace(std::numeric_limits<uint16_t>::max() - priority, impl_->task_seq, start_time, callback);
        ++impl_->task_seq;
    }

    void MessageLoop::do_consume()
    {
        std::unique_lock lock(impl_->mtx);
        impl_->is_running = true;
        impl_->is_busy = true;
        impl_->thread_id = std::this_thread::get_id();
        int64_t sleep_time = -1;
        decltype(impl_->task_queue) temp_queue;
        const auto &predicate_func = [this]() -> bool
        {
            return impl_->quit_flag || impl_->is_busy || !impl_->task_queue.empty();
        };
        on_begin();
        for (;;)
        {
            impl_->task_seq = 0;
            temp_queue.swap(impl_->task_queue);
            lock.unlock();
            while (!temp_queue.empty())
            {
                auto &&[priority, seq, start_time, task] = std::move(const_cast<MessageLoopImpl::TaskTuple &>(temp_queue.top()));
                on_task_changed(task, start_time);
                temp_queue.pop();
            }
            on_idle();
            if (impl_->quit_flag)
            {
                break;
            }
            lock.lock();
            if (!impl_->timer_set.empty())
            {
                process_timer_task(sleep_time);
            }
            impl_->is_busy = false;
            impl_->cv.notify_all();
            if (sleep_time < 0)
            {
                impl_->cv.wait(lock, predicate_func);
            }
            else if (sleep_time > 0)
            {
                impl_->cv.wait_for(lock, std::chrono::nanoseconds(sleep_time), predicate_func);
            }
            impl_->is_busy = true;
        }
        on_end();
        lock.lock();
        impl_->thread_id = std::thread::id();
        impl_->is_busy = false;
        impl_->is_running = false;
        impl_->cv.notify_all();
    }

    void MessageLoop::process_timer_task(int64_t &next_sleep_time)
    {
        uint64_t current_time = get_current_nano_time();
        uint64_t invoke_count = 0;
        int64_t interval_time = 0;
        int64_t remain_time = -1;
        next_sleep_time = -1;
        bool has_erase = false;
        for (auto iter = impl_->timer_set.begin(); iter != impl_->timer_set.end();)
        {
            Timer *timer = *iter;
            if (!timer->is_active())
            {
                ++iter;
                continue;
            }
            if (!timer->get_callback())
            {
                timer->stop();
                ++iter;
                continue;
            }
            interval_time = static_cast<uint64_t>(timer->get_interval()) * 1000'000U;
            if (interval_time <= 0)
            {
                interval_time = Timer::kMinInterval;
            }
            invoke_count = (current_time - timer->get_start_time() + Timer::kMinInterval) / interval_time;
            if (invoke_count > timer->get_invoke_count())
            {
                has_erase = false;
                for (uint64_t i = 0; timer->get_remain_loop_count() != 0 && i < invoke_count - timer->get_invoke_count(); ++i)
                {
                    if (impl_->task_queue.size() >= get_max_task_count())
                    {
                        LOG_WARN("Timer task is full(%s).", impl_->name.c_str());
                        std::this_thread::sleep_for(std::chrono::milliseconds(5));
                    }
                    emplace_task(timer->get_callback(), timer->get_priority());
                    if (timer->is_once_type())
                    {
                        delete timer;
                        iter = impl_->timer_set.erase(iter);
                        has_erase = true;
                        break;
                    }
                    timer->sub_remain_loop_count();
                }
                if (has_erase)
                {
                    continue;
                }
                if (timer->get_remain_loop_count() == 0)
                {
                    timer->stop();
                }
                else
                {
                    timer->set_invoke_count(invoke_count);
                }
                next_sleep_time = 0;
            }
            else
            {
                remain_time = interval_time - (current_time - timer->get_start_time()) % interval_time - Timer::kMinInterval;
                if (remain_time < 0)
                {
                    remain_time = 0;
                }
                if (next_sleep_time < 0 || next_sleep_time > remain_time)
                {
                    next_sleep_time = remain_time;
                }
            }
            ++iter;
        }
    }

} // namespace zeroplus::mq
